# -*- coding:utf-8 -*-
from . import app_register
from flask import render_template, request, flash, redirect, url_for
from app import db, Users
from werkzeug.security import generate_password_hash, check_password_hash
import pyotp
import os
import traceback
from qrcode import QRCode, constants
import base64


@app_register.route("/")
def index():
    return render_template('register.html')



@app_register.route("/register_write_mysql", methods=['POST'])
def register_write_mysql():
    if request.method == 'POST':
        username = request.form.get('username')
        passwd = request.form.get('passwd')
        passwd = generate_password_hash(passwd)  ##密码hash加密
        gtoken = pyotp.random_base32(64)  # 获取随机密钥，用于谷歌验证器,随机64位
        if not username or not passwd:
            flash('输入有误')
        else:
            users = Users.query.filter_by(username=username).all()  ##查询是否有用户，有就返回列表
            if len(users) == 0:
                sql = Users(username, passwd, gtoken)
                db.session.add(sql)
                db.session.commit()
                return redirect(url_for('app_register.mfa', gtoken=gtoken, username=username))
            else:
                flash('用户已存在')
    return redirect(url_for('app_register.index'))


@app_register.route("/mfa/?<string:gtoken>&?<string:username>")
def mfa(gtoken, username):
    ret = get_qrcode(gtoken, username)
    if ret[0]:
        img_stream = return_img_stream(ret[1])
    else:
        img_stream = None
    return render_template('register_mfa.html', img_stream=img_stream,gtoken=gtoken)


def read_static():
    print(app_register.root_path)
    # 读取文件 /home/bjhee/flask-app/admin/files/info.txt
    with app_register.open_resource('static/file/readtest.txt') as f:
        info = f.read()
    print(info)


def get_qrcode(secret_key, username):
    BASE_DIR = app_register.root_path
    filepath = BASE_DIR + os.sep + 'static' + os.sep + 'img' + os.sep + secret_key + '.png'
    data = pyotp.totp.TOTP(secret_key).provisioning_uri(username, issuer_name="MFA Code")
    qr = QRCode(
        version=1,
        error_correction=constants.ERROR_CORRECT_L,
        box_size=6,
        border=4, )
    try:
        qr.add_data(data)
        qr.make(fit=True)
        img = qr.make_image()
        img.save(filepath)  # 保存条形码图片
        return True, filepath
    except Exception:
        return False, None




def return_img_stream(img_local_path):
    """
    工具函数:
    获取本地图片流
    :param img_local_path:文件单张图片的本地绝对路径
    :return: 图片流
    """
    img_stream = ''
    with open(img_local_path, 'rb') as img_f:
        img_stream = img_f.read()
        img_stream = base64.b64encode(img_stream) #Base64是一种用64个字符来表示任意二进制数据的方法。
        img_stream=img_stream.decode()  ##bytes转成字符串
    return img_stream
